"""数据转换器，用于将清洗后的数据转换为模型训练所需的格式"""
import json
import numpy as np
from datetime import datetime

from fin_senti_entity_platform.utils import ENTITY_TYPES, SENTIMENT_LABELS
from fin_senti_entity_platform.utils.config_loader import config_loader
from fin_senti_entity_platform.utils.logger import Logger

# 初始化日志记录器
logger = Logger.get_logger('data_converter', 'data_processing.log')

class DataConverter:
    """
    数据转换器，负责将清洗后的数据转换为模型训练所需的格式
    """
    def __init__(self):
        """
        初始化数据转换器
        """
        # 加载转换配置
        self.convert_config = config_loader.get('data_collection', {}).get('data_processing', {}).get('conversion', {})
        
        # 实体类型映射
        self.entity_type_map = {entity: idx for idx, entity in enumerate(ENTITY_TYPES)}
        
        # 情感标签映射
        self.sentiment_label_map = {label: idx for idx, label in enumerate(SENTIMENT_LABELS)}
        
        # 最大序列长度
        self.max_seq_length = self.convert_config.get('max_seq_length', 512)
        
        logger.info("数据转换器初始化完成")
        
    def convert_to_model_input(self, cleaned_data, model_type='entity_recognition'):
        """
        将清洗后的数据转换为模型输入格式
        :param cleaned_data: 清洗后的数据
        :param model_type: 模型类型（entity_recognition或sentiment_analysis）
        :return: 模型输入数据
        """
        if not cleaned_data:
            return None
            
        try:
            if model_type == 'entity_recognition':
                return self._convert_for_entity_recognition(cleaned_data)
            elif model_type == 'sentiment_analysis':
                return self._convert_for_sentiment_analysis(cleaned_data)
            else:
                logger.error(f"不支持的模型类型: {model_type}")
                return None
                
        except Exception as e:
            logger.error(f"转换数据为模型输入失败: {str(e)}")
            return None
            
    def _convert_for_entity_recognition(self, cleaned_data):
        """
        转换为实体识别模型输入格式
        :param cleaned_data: 清洗后的数据
        :return: 实体识别模型输入数据
        """
        model_input = {
            'text_id': cleaned_data.get('text_id', str(datetime.now().timestamp())),
            'text': cleaned_data.get('normalized_content', ''),
            'tokens': cleaned_data.get('tokens', []),
            'keywords': cleaned_data.get('keywords', []),
            'meta_info': {
                'url': cleaned_data.get('url', ''),
                'title': cleaned_data.get('title', ''),
                'publish_time': cleaned_data.get('publish_time', ''),
                'source': cleaned_data.get('source', ''),
                'domain': cleaned_data.get('domain', '')
            },
            # 这里预留标签位置，实际标签会在标注阶段添加
            'labels': []
        }
        
        return model_input
        
    def _convert_for_sentiment_analysis(self, cleaned_data):
        """
        转换为情感分析模型输入格式
        :param cleaned_data: 清洗后的数据
        :return: 情感分析模型输入数据
        """
        model_input = {
            'text_id': cleaned_data.get('text_id', str(datetime.now().timestamp())),
            'text': cleaned_data.get('normalized_content', ''),
            'tokens': cleaned_data.get('tokens', []),
            'keywords': cleaned_data.get('keywords', []),
            'meta_info': {
                'url': cleaned_data.get('url', ''),
                'title': cleaned_data.get('title', ''),
                'publish_time': cleaned_data.get('publish_time', ''),
                'source': cleaned_data.get('source', ''),
                'domain': cleaned_data.get('domain', '')
            },
            # 这里预留标签位置，实际标签会在标注阶段添加
            'sentiment_label': None,
            'sentiment_score': None
        }
        
        return model_input
        
    def batch_convert(self, cleaned_data_list, model_type='entity_recognition'):
        """
        批量转换数据
        :param cleaned_data_list: 清洗后的数据列表
        :param model_type: 模型类型
        :return: 转换后的模型输入数据列表
        """
        if not cleaned_data_list or not isinstance(cleaned_data_list, list):
            return []
            
        converted_data = []
        
        for cleaned_data in cleaned_data_list:
            model_input = self.convert_to_model_input(cleaned_data, model_type)
            if model_input:
                converted_data.append(model_input)
                
        logger.info(f"批量转换完成: {len(converted_data)}/{len(cleaned_data_list)} 条数据成功转换")
        
        return converted_data
        
    def pad_sequences(self, sequences, max_length=None, padding='post', truncating='post', value=0):
        """
        对序列进行填充和截断
        :param sequences: 序列列表
        :param max_length: 最大长度
        :param padding: 填充方式（'pre'或'post'）
        :param truncating: 截断方式（'pre'或'post'）
        :param value: 填充值
        :return: 填充和截断后的序列数组
        """
        if max_length is None:
            max_length = self.max_seq_length
            
        # 计算实际的最大长度
        actual_max_length = min(max_length, max(len(seq) for seq in sequences) if sequences else 0)
        
        # 初始化结果数组
        result = np.full((len(sequences), actual_max_length), value, dtype=np.int32)
        
        # 填充序列
        for i, seq in enumerate(sequences):
            seq_len = len(seq)
            
            if seq_len == 0:
                continue
                
            if seq_len <= actual_max_length:
                # 序列长度小于最大长度，进行填充
                if padding == 'post':
                    result[i, :seq_len] = seq
                else:
                    result[i, -seq_len:] = seq
            else:
                # 序列长度大于最大长度，进行截断
                if truncating == 'post':
                    result[i, :] = seq[:actual_max_length]
                else:
                    result[i, :] = seq[-actual_max_length:]
                    
        return result
        
    def convert_labels_to_ids(self, labels, label_map):
        """
        将标签转换为ID
        :param labels: 标签列表
        :param label_map: 标签映射字典
        :return: 标签ID列表
        """
        label_ids = []
        
        for label in labels:
            label_ids.append(label_map.get(label, -1))  # 未知标签用-1表示
            
        return label_ids
        
    def convert_ids_to_labels(self, label_ids, id_map):
        """
        将ID转换为标签
        :param label_ids: 标签ID列表
        :param id_map: ID映射字典
        :return: 标签列表
        """
        labels = []
        
        for label_id in label_ids:
            labels.append(id_map.get(label_id, 'UNKNOWN'))
            
        return labels
        
    def save_converted_data(self, converted_data, output_file):
        """
        保存转换后的数据到文件
        :param converted_data: 转换后的数据
        :param output_file: 输出文件路径
        :return: 是否保存成功
        """
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                for item in converted_data:
                    f.write(json.dumps(item, ensure_ascii=False) + '\n')
                    
            logger.info(f"转换后的数据已保存到: {output_file}")
            return True
            
        except Exception as e:
            logger.error(f"保存转换后的数据失败: {str(e)}")
            return False
            
    def load_converted_data(self, input_file):
        """
        从文件加载转换后的数据
        :param input_file: 输入文件路径
        :return: 转换后的数据列表
        """
        converted_data = []
        
        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                for line in f:
                    try:
                        item = json.loads(line.strip())
                        converted_data.append(item)
                    except Exception:
                        continue
                        
            logger.info(f"已从 {input_file} 加载 {len(converted_data)} 条转换后的数据")
            return converted_data
            
        except Exception as e:
            logger.error(f"加载转换后的数据失败: {str(e)}")
            return []
